#include <arpa/inet.h>
#include <cstring>
#include <fcntl.h>
#include <iostream>
#include <netinet/in.h>
#include <sys/socket.h>
#include <unistd.h>

// 定义错误提示宏
#define ERROR_MSG(fmt, ...) printf("Error: " fmt "\n", ##__VA_ARGS__)

// 定义协议枚举
enum class Protocol { TCP = IPPROTO_TCP, UDP = IPPROTO_UDP };

// 创建套接字
int createSocket(Protocol protocol) {
	// 根据协议类型创建相应的套接字类型
	int type   = (protocol == Protocol::TCP) ? SOCK_STREAM : SOCK_DGRAM;
	int sockfd = socket(AF_INET, type, static_cast<int>(protocol)); // AF_INET表示IPv4协议
	if (sockfd == -1) {
		ERROR_MSG("Failed to create socket");
	}

	// 设置地址复用选项，避免地址被占用时无法绑定
	int reuse = 1;
	if (setsockopt(sockfd, SOL_SOCKET, SO_REUSEADDR, &reuse, sizeof(reuse)) == -1) {
		ERROR_MSG("Failed to set SO_REUSEADDR option");
		close(sockfd);
		return -1;
	}

	return sockfd;
}

// 创建 sockaddr 地址
void createAddress(const char *ip, int port, struct sockaddr_in &addr) {
	memset(&addr, 0, sizeof(addr));
	addr.sin_family = AF_INET;
	addr.sin_port   = htons(port);
	// 将点分十进制的 IP 地址转换为二进制形式
	inet_pton(AF_INET, ip, &addr.sin_addr);
}

// 绑定地址
bool bindAddress(int sockfd, const struct sockaddr_in &addr) {
	if (bind(sockfd, (const struct sockaddr *) &addr, sizeof(addr)) == -1) {
		ERROR_MSG("Failed to bind address");
		return false;
	}
	return true;
}

// 监听
bool listenOnSocket(int sockfd, int backlog) {
	if (listen(sockfd, backlog) == -1) {
		ERROR_MSG("Failed to listen on socket");
		return false;
	}
	return true;
}

// 使用 select 监听
bool selectListen(int sockfd, int timeout_sec) {
	fd_set readfds;
	struct timeval timeout;
	timeout.tv_sec  = timeout_sec;
	timeout.tv_usec = 0;

	FD_ZERO(&readfds);
	FD_SET(sockfd, &readfds);

	int activity = select(sockfd + 1, &readfds, NULL, NULL, &timeout);
	if (activity == -1) {
		ERROR_MSG("Select error");
		return false;
	}
	return activity > 0;
}

// 创建组播
int createMulticast(const char *multicast_ip, int port) {
	// 创建 UDP 套接字用于组播
	int sockfd = createSocket(Protocol::UDP);
	if (sockfd == -1) {
		return -1;
	}

	struct sockaddr_in addr;
	createAddress(multicast_ip, port, addr);

	struct ip_mreq mreq;
	// 设置组播组的 IP 地址
	mreq.imr_multiaddr.s_addr = inet_addr(multicast_ip);
	// 设置本地接口地址，使用 INADDR_ANY 表示使用任意可用接口
	mreq.imr_interface.s_addr = htonl(INADDR_ANY);

	// 设置加入组播组的选项
	if (setsockopt(sockfd, IPPROTO_IP, IP_ADD_MEMBERSHIP, &mreq, sizeof(mreq)) == -1) {
		ERROR_MSG("Failed to set multicast options");
		close(sockfd);
		return -1;
	}

	if (!bindAddress(sockfd, addr)) {
		close(sockfd);
		return -1;
	}

	return sockfd;
}
